import datasets
from datasets import load_dataset
from unsloth import FastLanguageModel, is_bfloat16_supported

# 导入必要的库和模块
import os
from datasets import load_dataset

# 设置环境变量以使用镜像站点
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

# 定义模型参数
max_seq_length = 256
lora_rank = 16
model, tokenizer = FastLanguageModel.from_pretrained(
    "unsloth/DeepSeek-R1-Distill-Llama-8B",
    max_seq_length=max_seq_length,
    max_lora_rank=16,
    load_in_4bit=True,  # 使用4-bit量化加载模型以节省显存
    fast_inference=True,  # 启用快速推理模式
    gpu_memory_utilization=0.9,  # 设置GPU显存利用率
    local_files_only=True,  # 仅使用本地文件
)
prompt_style = """以下是描述任务的指令，以及提供更多上下文的输入。
请写出恰当完成该请求的回答。
在回答之前，请仔细思考问题，以确保回答合乎逻辑且准确。
### Instruction:
你是一位在临床推理、诊断和治疗计划方面具有专业知识的医学专家。
请回答以下医学问题。
### description:
{}
### question:
{}
### answer:
{}"""

# 测试用医学问题
desc = "一名70岁的男性患者因胸痛伴呕吐16小时就医，心电图显示下壁导联和右胸导联ST段抬高0.1~0.3mV，经补液后血压降至80/60mmHg，患者出现呼吸困难和不能平卧的症状，体检发现双肺有大量水泡音。"
question = "在这种情况下，最恰当的药物处理是什么？"
# 设置模型为推理模式并生成回答
FastLanguageModel.for_inference(model)
inputs = tokenizer([prompt_style.format(desc, question, "")], return_tensors="pt").to(
    "cuda"
)
outputs = model.generate(
    input_ids=inputs.input_ids,  # 输入token的id序列
    attention_mask=inputs.attention_mask,  # 注意力掩码，标记有效输入位置
    max_new_tokens=1024,  # 生成的最大新token数量
    use_cache=True,  # 使用KV缓存加速生成
)
response = tokenizer.batch_decode(outputs)
print("### 微调前模型推理结果：")
print(response)


# 在数据加载后，展开 QA_pairs 为独立样本
def explode_qa_pairs(examples):
    exploded_data = {"title": [], "description": [], "question": [], "answer": []}
    title = examples["title"]
    desc = examples["description"]
    qa_pairs = examples["QA_pairs"]
    for pair in qa_pairs:
        exploded_data["title"].append(title)
        exploded_data["description"].append(desc)
        exploded_data["question"].append(pair["question"])
        exploded_data["answer"].append(pair["answer"])
    return exploded_data


def formatting_prompts_func(examples):
    texts = []
    for title, desc, question, answer in zip(
        examples["title"],
        examples["description"],
        examples["question"],
        examples["answer"],
    ):
        text = prompt_style.format(desc, question, answer) + tokenizer.eos_token
        texts.append(text)
    return {"text": texts}


dataset = load_dataset(
    "json",
    data_files="/home/patrick/hub/datasets--FreedomIntelligence--CMB/snapshots/935fbc09edf1303d89872b21265ff597f426ac0d/CMB-Clin-qa.json",
    trust_remote_code=True,
    split="train",
)

print(dataset.column_names)  # 打印数据集列名

# 展开数据
dataset = dataset.map(explode_qa_pairs, batched=False)
dataset = dataset.remove_columns(["QA_pairs"])  # 移除原始嵌套字段
all_dataset = dataset.train_test_split(test_size=0.01)
train_data = all_dataset["train"]
test_data = all_dataset["test"]
# 重新执行数据映射
train_data = train_data.map(formatting_prompts_func, batched=True)
test_data = test_data.map(formatting_prompts_func, batched=True)
# 配置LoRA微调模型
model = FastLanguageModel.get_peft_model(
    model,
    r=16,  # LoRA秩，建议值为8, 16, 32等
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ],  # 目标模块
    lora_alpha=16,  # LoRA缩放因子
    lora_dropout=0,  # Dropout设置为0以优化性能
    bias="none",  # 偏置设置为"none"以优化性能
    use_gradient_checkpointing="unsloth",  # 使用梯度检查点以节省显存
    random_state=3407,  # 随机种子
    use_rslora=False,  # 不使用rank stabilized LoRA
    loftq_config=None,  # 不使用LoftQ配置
)
# print(model)  # 打印模型结构

# 配置训练器
from trl import SFTTrainer, GRPOTrainer
from transformers import TrainingArguments

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=train_data,
    eval_dataset=test_data,
    dataset_text_field="text",  # 数据集中文本字段
    max_seq_length=max_seq_length,
    dataset_num_proc=2,  # 并行处理进程数
    packing=False,  # 不启用packing以加快短序列训练
    args=TrainingArguments(
        per_device_train_batch_size=4,  # 每个设备的批量大小
        gradient_accumulation_steps=4,  # 梯度累积步数
        eval_strategy="steps",
        eval_steps=10,
        warmup_steps=5,  # 预热步数
        max_steps=50,  # 最大训练步数
        learning_rate=2e-5,  # 学习率
        fp16=not is_bfloat16_supported(),  # 根据硬件支持选择精度
        bf16=is_bfloat16_supported(),
        logging_steps=10,  # 日志记录步数
        optim="adamw_8bit",  # 使用8-bit AdamW优化器
        weight_decay=0.1,  # 权重衰减
        lr_scheduler_type="cosine",  # 学习率调度器类型
        seed=3407,  # 随机种子
        output_dir="outputs",  # 输出目录
        report_to="tensorboard",  # 不报告到外部工具
        logging_dir="./logs",
    ),
)

# 开始训练

print(f"trainer.args.max_steps: {trainer.args.max_steps}")
print(f"trainer.args.num_train_epochs: {trainer.args.num_train_epochs}")

print(f"trainer.args.num_train_epochs: {trainer.args.num_train_epochs}")
print(f"trainer.args.eval_steps: {trainer.args.eval_steps}")
trainer.train()
print(f"Total training steps: {trainer.state.max_steps}")
print(f"Total epochs: {trainer.state.epoch}")

# 微调后再次生成回答
FastLanguageModel.for_inference(model)
inputs = tokenizer([prompt_style.format(desc, question, "")], return_tensors="pt").to(
    "cuda"
)
outputs = model.generate(
    input_ids=inputs.input_ids,
    attention_mask=inputs.attention_mask,
    max_new_tokens=1024,
    use_cache=True,
)
response = tokenizer.batch_decode(outputs)
print("### 微调后模型推理结果：")
print(response)
model.save_pretrained("dsk_med_lora_model")  # Local saving
tokenizer.save_pretrained("dsk_med_lora_model")
